#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复attach_2_file表中location为空的记录
"""

import os
import sys
import uuid
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from db.session import get_session
from db.models.attachments import Attach2File
from db.models.files import QAFile
from loguru import logger
from utils.file_utils import download_attachment


def find_missing_locations():
    """查找location为空或文件不存在的记录"""
    with get_session() as session:
        # 查找location为空或文件不存在的记录
        records = session.query(Attach2File).filter(
            Attach2File.download_link.isnot(None),
            Attach2File.download_link != '',
            Attach2File.is_delete == False
        ).all()
        
        missing_records = []
        for record in records:
            if not record.location or not os.path.exists(record.location):
                missing_records.append(record)
        
        logger.info(f"找到 {len(missing_records)} 条需要修复的记录")
        return missing_records


def fix_attach_record(record: Attach2File, save_dir: str):
    """修复单条记录"""
    try:
        if not record.download_link:
            logger.warning(f"记录 {record.id} 没有download_link，跳过")
            return False
        
        # 获取文件对象以确定保存目录
        file_obj = record.file
        if not file_obj:
            logger.warning(f"记录 {record.id} 没有关联的文件对象，跳过")
            return False
        
        # 构建保存目录
        file_dir = os.path.dirname(file_obj.location)
        if not file_dir:
            logger.warning(f"文件 {file_obj.id} 没有有效的location，跳过")
            return False
        
        # 尝试重新下载
        logger.info(f"尝试重新下载: {record.download_link}")
        status, img_path = download_attachment(record.download_link, file_dir)
        
        if status and os.path.exists(img_path):
            # 更新记录
            record.location = img_path
            record.origin_name = os.path.basename(record.download_link)
            
            with get_session() as session:
                session.merge(record)
                session.commit()
            
            logger.info(f"成功修复记录 {record.id}: {img_path}")
            return True
        else:
            logger.warning(f"下载失败: {record.download_link}")
            return False
            
    except Exception as e:
        logger.error(f"修复记录 {record.id} 时发生错误: {e}")
        return False


def main():
    """主函数"""
    logger.info("开始修复attach_2_file表中的记录...")
    
    # 查找需要修复的记录
    missing_records = find_missing_locations()
    
    if not missing_records:
        logger.info("没有需要修复的记录")
        return
    
    # 修复记录
    success_count = 0
    for record in missing_records:
        if fix_attach_record(record, ""):
            success_count += 1
    
    logger.info(f"修复完成: {success_count}/{len(missing_records)} 条记录修复成功")


if __name__ == "__main__":
    main()
